import random
import json
from .openai_model import OpenAIModel
import os

from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated
import operator
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, ToolMessage
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import InMemorySaver

# 用于存储智能体的消息
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]

# 增加了记忆参数 checkpointer
class Agent:
    def __init__(self, model, tools, checkpointer, system=""):
        self.system = system
        graph = StateGraph(AgentState)
        graph.add_node("llm", self.call_openai)
        graph.add_node("action", self.take_action)
        graph.add_conditional_edges("llm", self.exists_action, {True: "action", False: END})
        graph.add_edge("action", "llm")
        graph.set_entry_point("llm")

        # 编译图的时候传入 checkpointer
        self.graph = graph.compile(checkpointer=checkpointer)
        
        self.tools = {t.name: t for t in tools} if tools else None
        self.model = model.bind_tools(tools) if tools else model


    def exists_action(self, state: AgentState):
        """
        判断是否需要调用工具，从 LLM 返回的消息中判断是否存在 tool_calls 字段，如果有则返回 True
        注意：硅基流动的 deepseek 模型不支持 tool_calls 字段，Qwen3 支持
        """
        result = state['messages'][-1]
        return len(result.tool_calls) > 0

    def call_openai(self, state: AgentState):
        """
        调用模型，传入消息列表。如果设置了系统提示词，则先添加到消息列表中
        """
        messages = state['messages']
        print("消息条数：{}".format(len(messages)))
        if len(messages) ==1 and self.system:
            print("添加系统提示词")
            messages = [SystemMessage(content=self.system)] + messages
        message = self.model.invoke(messages)
        return {'messages': [message]}

    def take_action(self, state: AgentState):
        """
        根据 LLM 返回的消息中的 tool_calls 字段调用对应的工具
        """
        tool_calls = state['messages'][-1].tool_calls
        # print(tool_calls)
        results = []
        for t in tool_calls:
            print(f"调用工具: {t}")
            if not t['name'] in self.tools:      # 是否存在这个工具
                print("\n ....未知工具....")
                result = "未知工具名称，请重试"  
            else:
                # 调用工具
                result = self.tools[t['name']].invoke(t['args'])
                print(f"工具返回结果： {result}")
            # 将工具返回的结果放入返回的消息列表中
            results.append(ToolMessage(tool_call_id=t['id'], name=t['name'], content=str(result)))
        print("返回大模型继续对话!")
        return {'messages': results}

class QuestionGenerator:
    """题目生成器类"""
    def __init__(self, knowledge_manager, session_id=None):
        self.knowledge_manager = knowledge_manager
        self.ai_model = OpenAIModel()
        self.session_id = session_id # langgraph used for session
        
        self.thread = {"configurable": {"thread_id": str(session_id)}} 
        self.memory = InMemorySaver()

        self.model = ChatOpenAI(
            base_url=os.getenv("API_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1"),
            api_key= os.getenv("API_KEY", "API_KEY"),
            model=os.getenv("MODEL_NAME", "qwen3-32b"),
            extra_body={"enable_thinking": False}  # 百炼 qwen3-32b 思考类模型需要增加参数
        )  
        self.agent = Agent(
            model=self.model,
            tools=None,
            checkpointer=self.memory,
            system="你是一个专业的智能体，用于生成基于以下知识点的选择题题目，生成题目不能重复，可以逐渐增加难度。"
        )
    
    def generate_question(self, knowledge_point_id, selected_kps):
        """生成题目"""
        # 使用AI模型生成题目
        # ai_question = self._generate_with_ai(knowledge_point_id)
        # 使用LangGraph生成题目
        lg_question = self._generate_with_lg(knowledge_point_id)
        return lg_question
    
    # 获取知识点信息
    def _get_knowledge(self, knowledge_point_id):
        # 获取知识点名称
        kp_name = "未知知识点"
        for grade in self.knowledge_manager.knowledge_points.values():
            for subject in grade.values():
                for kp in subject:
                    if kp["id"] == knowledge_point_id:
                        kp_name = kp["name"]
                        # 判断科目类型
                        subject_type = "math" if "math" in kp["id"] else "chinese"
                        break
                else:
                    continue
                break
            else:
                continue
            break
        return kp_name, subject_type

    def _generate_with_ai(self, knowledge_point_id):
        """使用AI模型生成题目"""
        # 获取知识点
        kp_name, subject_type = self._get_knowledge(knowledge_point_id)
        if kp_name == "未知知识点":
            return None
        
        # 构建提示词
        if subject_type == "math":
            prompt = f"为小学生生成一道关于'{kp_name}'的数学选择题，包含4个选项（A、B、C、D），其中只有一个正确答案。请以JSON格式返回，包含text（题目内容）、options（选项字典，键为A/B/C/D，值为选项内容）、correct_answer（正确选项）、explanation（答案解释）。"
        else:
            prompt = f"为小学生生成一道关于'{kp_name}'的语文选择题，包含4个选项（A、B、C、D），其中只有一个正确答案。请以JSON格式返回，包含text（题目内容）、options（选项字典，键为A/B/C/D，值为选项内容）、correct_answer（正确选项）、explanation（答案解释）。"
        
        # 调用AI模型
        try:
            response = self.ai_model.generate_text(prompt)
            if response:
                # 解析JSON响应
                if response.startswith("```json") and response.endswith("```"):
                    response = response[7:-3].strip()
                
                result = json.loads(response)
                result["knowledge_point_id"] = knowledge_point_id
                return result
        except Exception as e:
            print(f"AI生成题目失败: {e}")
        
        return None
    
    def _generate_with_lg(self, knowledge_point_id):
        """使用LangGraph生成题目"""
        # 获取知识点
        kp_name, subject_type = self._get_knowledge(knowledge_point_id)
        if kp_name == "未知知识点":
            return None
        
        # 构建提示词
        prompt = f"为小学生生成一道关于'{kp_name}'的题目，包含4个选项（A、B、C、D），其中只有一个正确答案。请以JSON格式返回，包含text（题目内容）、options（选项字典，键为A/B/C/D，值为选项内容）、correct_answer（正确选项）、explanation（答案解释）。"
        messages = [HumanMessage(content=prompt)]
        # 调用LangGraph
        result = self.agent.graph.invoke({"messages": messages}, config=self.thread)
        # 解析JSON响应
        if result["messages"][-1].content.startswith("```json") and result["messages"][-1].content.endswith("```"):
            result["messages"][-1].content = result["messages"][-1].content[7:-3].strip()
        result = json.loads(result["messages"][-1].content)
        result["knowledge_point_id"] = knowledge_point_id
        return result

    def select_knowledge_point(self, selected_kps, learning_memory):
        """根据学习情况选择知识点"""
        # 如果只有一个知识点，直接返回
        if len(selected_kps) == 1:
            return selected_kps[0]
        
        # 计算每个知识点的正确率
        kp_accuracies = {kp_id: learning_memory.get_accuracy(kp_id) for kp_id in selected_kps}
        
        # 按正确率排序，优先选择正确率低的知识点
        sorted_kps = sorted(kp_accuracies.items(), key=lambda x: x[1])
        
        # 基于正确率分配权重，正确率越低的权重越高
        weights = [max(0.1, 1 - acc) for _, acc in sorted_kps]
        total_weight = sum(weights)
        normalized_weights = [w/total_weight for w in weights]
        
        # 根据权重随机选择一个知识点
        r = random.random()
        cumulative_weight = 0
        for i, weight in enumerate(normalized_weights):
            cumulative_weight += weight
            if r <= cumulative_weight:
                return sorted_kps[i][0]
        
        # 兜底返回第一个知识点
        return selected_kps[0]